import angr
import sys

import angr.knowledge_base

sys.path.append('src')

from tasks.lsbvd.feature.const_value import buffer_overflow_funcs,buffer_overflow_sink_funcs,buffer_overflow_source_funcs


class CallRelationShipFeatures:
    def __init__(self, cfg_fast:angr.analyses.cfg.cfg_fast.CFGFast):
        self.cfg= cfg_fast

    def get_func_calls(self, func_addrs):
        """
        获取函数调用点
        :param func_addrs: 函数地址列表
        :return: 函数调用点列表
        """
        return get_func_calls(self.cfg, func_addrs)

class SinkOrientSearcher:
    def __init__(self,proj:angr.Project, cfg_fast:angr.analyses.cfg.cfg_fast.CFGFast,name2addr:dict=None):
        self.cfg = cfg_fast
        self.proj = proj
        self.name2addr=name2addr
        self.d0_sink=self._get_d0_list(buffer_overflow_sink_funcs())
        self.d0_source=self._get_d0_list(buffer_overflow_source_funcs())
        self.depth_caches_source={0:self.d0_source}
        self.depth_caches_sink={0:self.d0_sink}

    def source_sink_pair(self):
        source_paths=[]
        for source in self.d0_source:
            source_paths+=self.call_path(source.addr)
        sink_paths=[]
        for sink in self.d0_sink:
            sink_paths+=self.call_path(sink.addr)
        for source_path in source_paths:
            for sink_path in sink_paths:
                common_nodes=self.common_nodes(source_path,sink_path)
                if len(common_nodes)>0:
                    yield source_path,sink_path,common_nodes

    def generate_source_paths(self):
        """
        生成源函数调用路径
        :return: 源函数调用路径列表
        """
        for source in self.d0_source:
            paths = self.call_path(source.addr)
            for path in paths:
                yield path

    def generate_sink_paths(self):
        """
        生成汇函数调用路径
        :return: 汇函数调用路径列表
        """
        for sink in self.d0_sink:
            paths = self.call_path(sink.addr)
            for path in paths:
                yield path

    def _get_d0_list(self,func_names):
        candidate_names=func_names
        ret=[]
        for name in candidate_names:
            func = self.proj.kb.functions.function(addr=self.name2addr[name])
            if func is not None:
                ret.append(func)
        return ret

    def depth_n_sink(self,n:int):
        return self.depth_n(n,self.depth_caches_sink)
    
    def depth_n_source(self,n:int):
        return self.depth_n(n,self.depth_caches_source)

    def depth_n(self,n:int,caches):
        """
        获取深度为n的函数
        :param n: 深度
        :return: 函数列表
        """
        if n in caches:
            return caches[n]
        previous_depth=self.depth_n(n-1,caches)
        ret=set()
        for func in previous_depth:
            try:
                predecessors = self.call_graph.predecessors(func.addr)
                for predecessor in predecessors:
                    if predecessor not in previous_depth:
                        ret.add(self.proj.kb.functions.function(addr=predecessor))
            except:
                continue
        self.depth_caches_sink[n]=list(ret)
        return ret

    def call_path(self, func_addr:int,max_number=1000000):
        # 以call_site为起点,深度遍历call_graph
        all_paths=[]
        def get_path(node,path:list,visited:set):
            visited.add(node)
            path.append(node)
            if len(all_paths) >= max_number:
                return
            if node not in self.call_graph:
                all_paths.append(path.copy())
                path.pop()
                visited.remove(node)
                return
            preds = list(self.call_graph.predecessors(node))
            if len(preds) == 0:
                all_paths.append(path.copy())
            else:
                for pred in preds:
                    if pred not in visited:
                        get_path(pred, path, visited)
            path.pop()
            visited.remove(node)
        get_path(func_addr,[],set())
        return all_paths

    def common_nodes(self, path1, path2):
        common_nodes = set(path1) & set(path2)
        return list(common_nodes)
    
    @property
    def func_number(self):
        return len(self.cfg.kb.functions)

    @property
    def call_graph(self):
        return self.cfg.kb.functions.callgraph

    def all_funcs(self):
        return self.cfg.kb.functions.values()


if __name__ == '__main__':
    import time
    import json
    import os
    import pickle
    from tasks.lsbvd.feature.static_cfg import BigBinaryCfgBuilder

    region_split_path='caches_data/amd.txtbin_partitions.json'
    region_json=json.load(open(region_split_path, 'r'))['partitions']

    bbinary_builder= BigBinaryCfgBuilder(
        binary_path='caches_data/ine.txtbin',
        caches_dir='caches_data/nokia_cfg',
        symbols_json_path='caches_data/symbols_export.json'
    )

    cfgs=bbinary_builder.load_partial_cfg(regions=region_json)
    cfg=None
    for e in cfgs:
        cfg=e
        break
    # 保存cfg
    
    cfg_save_path = 'caches_data/nokia_cfg/saved_cfg.pkl'
    os.makedirs(os.path.dirname(cfg_save_path), exist_ok=True)
    print(f"正在保存 CFG 到文件: {cfg_save_path}")
    start_time = time.time()
    with open(cfg_save_path, 'wb') as f:
        pickle.dump(cfg, f)
    print(f"CFG 保存完成，耗时: {time.time() - start_time:.2f} 秒")

    print("process cfg.")

    # proj=bbinary_builder.proj

    # searcher=SinkOrientSearcher(bbinary_builder.proj,cfg,bbinary_builder.name2iddr)

    # source_counter=0
    # sink_counter=0

    # with open('caches_data/nokia_path/source_path.txt','w') as f:
    #     for path in searcher.generate_source_paths():
    #         source_counter+=1
    #         if source_counter % 1000 == 0:
    #             print(f'Processed {source_counter} source paths')
    #         f.write(','.join([str(addr) for addr in path])+'\n')

    # with open('caches_data/nokia_path/sink_path.txt','w') as f:
    #     for path in searcher.generate_sink_paths():
    #         sink_counter+=1
    #         if sink_counter % 1000 == 0:
    #             print(f'Processed {sink_counter} sink paths')
    #         f.write(' '.join([str(addr) for addr in path])+'\n')

    # # ========== 新增：从指定地址向下遍历所有路径并保存到文件 ==========
    # def find_all_downward_paths(call_graph, start_addr, max_depth=20, max_paths=1000000):
    #     """
    #     从start_addr向下遍历所有路径，返回路径列表（每条路径为地址列表）
    #     """
    #     all_paths = []
    #     def dfs(node, path, visited, depth):
    #         if len(all_paths) >= max_paths:
    #             return
    #         path.append(node)
    #         visited.add(node)
    #         try:
    #             succs = list(call_graph.successors(node))
    #         except Exception:
    #             succs = []
    #         if not succs or depth >= max_depth:
    #             all_paths.append(path.copy())
    #         else:
    #             for succ in succs:
    #                 if succ not in visited:
    #                     dfs(succ, path, visited, depth + 1)
    #         path.pop()
    #         visited.remove(node)
    #     dfs(start_addr, [], set(), 0)
    #     return all_paths

    # def save_paths_to_file(paths, output_file):
    #     """
    #     将路径列表保存到文件，每行一条路径（地址用逗号分隔）
    #     """
    #     with open(output_file, 'w') as f:
    #         for path in paths:
    #             f.write(','.join(str(addr) for addr in path) + '\n')

    # # 目标函数地址（10进制）
    # # target_addrs = [60634032, 60634340, 60635240, 60637004]
    # target_addrs = [60635240, 60637004]
    # output_files = [
    # #    'caches_data/nokia_path/downward_paths_60634032.txt',
    # #    'caches_data/nokia_path/downward_paths_60634340.txt',
    #     'caches_data/nokia_path/downward_paths_60635240.txt',
    #     'caches_data/nokia_path/downward_paths_60637004.txt'
    # ]

    # for addr, out_file in zip(target_addrs, output_files):
    #     print(f"Generating all downward paths from function at address {addr} (hex: {hex(addr)}) ...")
    #     paths = find_all_downward_paths(cfg.kb.functions.callgraph, addr, max_depth=20, max_paths=1000000)
    #     print(f"Total downward paths from {hex(addr)}: {len(paths)}")
    #     save_paths_to_file(paths, out_file)
    #     print(f"Paths saved to {out_file}")